"""
Date: 2025/10/4
Author: 溪海莘
Project: TableTranslate
Explain: 智能表格翻译

1. 读取表格对应的列 - 并将其输出到text文件中
2. 构造Prompt内容 - 获取大模型的API - 对文本内容进行处理
3. 将大模型按照某一个格式返回的内容保存到csv文件中
4. 读取文本内容并写入表格对应的列 - translate

"""

import os
import pandas as pd
from typing import Optional
from openai import OpenAI
from dotenv import load_dotenv

# 加载 .env 文件中的环境变量
load_dotenv()  

# 配置API信息
API_KEY = os.environ.get("OPENAI_API_KEY")
API_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"  # API基础地址 - 火山引擎
MODEL_ID = "deepseek-r1-250120"     # 所用模型的ID
TEMP_FILE_READ_OUT = "temp_file.txt"
TEMP_FILE_RESPONSE_OUT = "AI_response.csv"
TRANSLATED_FILE = "translated_file.csv"


# 1. 读取表格对应的列
def read_table(file_path: str, sheetname: Optional[str] =None , skiprows: Optional[int] = 0, ) -> str:
    
    df = pd.read_excel(file_path, sheet_name=sheetname, 
                       skiprows=skiprows)    
    return df

def save_to_text(select_column: str, df, file_path: str):
    if select_column not in df.columns:
        print(f"File:{file_path} not exist...")
        return 
    
    data = df[select_column].astype(str)

    # if os.path.exists(TEMP_FILE_READ_OUT):
    #     return data
    
    with open(TEMP_FILE_READ_OUT, 'w', encoding='utf-8') as f:
        for item in data:
            f.write(f"{item}\n")
    return data

def read_data():
    with open(TEMP_FILE_READ_OUT, 'r') as f:
        data = f.readlines()
    return data

# 2. 构造Prompt内容 - 获取大模型的API - 对文本内容进行处理
def initialize_api_client() -> OpenAI:
    """初始化并返回OpenAI客户端"""
    return OpenAI(
        api_key=API_KEY,
        base_url=API_BASE_URL
    )

def prompt_excute(text_data, example_template, language):
    prompt = f"""
    任务: 翻译任务
    角色: 一名优秀的翻译官

    这是一些示例说明:
    {example_template}
    
    请结合示例说明为我翻译下面用【】包裹的内容, 将其翻译为{language}, 并按照以下格式返回: 
    第一个是原文, 第二个是译文, 两者中间以“@”分隔. 返回的译文无需添加额外的说明.
    并且一条内容对应一行.

    【{text_data}】

    """
    return prompt

def table_translate(client: OpenAI, prompt) -> str:
    try:
        # 调用大模型API
        response = client.chat.completions.create(
            model=MODEL_ID,
            messages=[
                {"role": "system", "content": "一名优秀的翻译官, 能将语句或单词精准地翻译为其它语言"},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,  # 降低随机性，使输出更稳定
            max_tokens=1500
        )
        
        # 提取API返回内容
        content = response.choices[0].message.content

        return content
    except Exception as e:
        print(f"生成解决方案时出错: {e}")
        
def add_translate_content(translated_df, df, select_column):
    df[select_column] = df[select_column].astype(str).str.strip()
    merged_df = pd.merge(df, translated_df, left_on=select_column, right_on='origin', how='left')
    merged_df.to_csv(TRANSLATED_FILE)

if __name__ == '__main__':
    TOKEN_LEN = 50
    test_df = read_table("./Data/test.xlsx")
    save_to_text(select_column='Country Name', file_path="./Data/test.xlsx", df=test_df['Sheet1'])
    text_data = read_data()
    print(text_data)
    print(type(text_data))  # list

    openai = initialize_api_client()
    example_template = """
                    China@中国
                    Australia@澳大利亚
                    Japan@日本
                    """
    cnt = []
    for i in range(0, len(text_data), TOKEN_LEN):
       cnt.append(i)    # 0, 50, 100, 150

    for j in range(len(cnt)):
        if j != len(cnt) -1:
            cnt_data = text_data[cnt[j]:cnt[j+1]]   # [0, 49], [50, 99]
        else:
            cnt_data = text_data[cnt[j]:]

        prompt = prompt_excute(cnt_data, example_template, '中文')
        content = table_translate(openai, prompt)
        print(content)

        with open(TEMP_FILE_RESPONSE_OUT, 'a', encoding='utf-8') as f:
            f.write(content + "\n")

    data = []
    with open(TEMP_FILE_RESPONSE_OUT, 'r', encoding='utf-8') as f:
        for cont in f.readlines():
            cont_striped = cont.strip().replace("'",'').replace("\\n", '')
            if not cont_striped :
                continue

            parts = cont_striped.split("@", 1)
            if len(parts) == 2:
                origin, translated = parts
                data.append([origin, translated])
            else:
                print(f"{cont_striped} translated faild!")
    df = pd.DataFrame(data, columns = ['origin', 'translated'])

    print(df)
    # 在调用函数前添加类型检查
    print(type(df))       # 应该输出 <class 'pandas.core.frame.DataFrame'>
    print(type(test_df['Sheet1']))  # 应该输出 <class 'pandas.core.frame.DataFrame'>
    if isinstance(test_df['Sheet1'], dict):
        test_df = pd.DataFrame(test_df['Sheet1'])
        print(test_df)
    add_translate_content(df, test_df['Sheet1'], 'Country Name')
    # add_translate_content(content, test_df, 'Country Name')
   






